#include "mt3d.h"

void mt3d::assemble_linear_system()
{
    std::clog << "Assembling kernel matrix ... ";
    size_t e, i, j;

    std::vector<triplt_cd> kernel_entries;
    kernel_entries.reserve(edge_num_*5);

    g3d_tet_.initiate(2);

    // 计算线性方程组的左端项
    A_PDE_.resize(edge_num_, edge_num_);
    A_PDE_.setZero();

    std::vector<std::vector<triplt_cd>> omp_entries;
    omp_entries.resize(36);
    for (i = 0; i < 36; i++)
    {
        omp_entries[i].resize(elem_num_);
    }

    edge_setup e_set;
    double edge_cndt, edge_cndt_ref, aij_p1, aij_p2;
    complex_d tmp_entry;
    gctl::vertex3dc* tmp_vert[4];
    // Loop all elements
#pragma omp parallel for private (e, i, j, e_set, edge_cndt, aij_p1, aij_p2, tmp_entry, tmp_vert) schedule(guided)
    for (e = 0; e < elem_num_; e++)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);

        for (i = 0; i < 6; i++)
        {
            e_set.uv_id[0] = i;

            if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
            else e_set.uv_type[0] = efem_sf::Normal;

            for (j = 0; j < 6; j++)
            {
                e_set.uv_id[1] = j;

                if (elem_edge_reversed_[e][j]) e_set.uv_type[1] = efem_sf::Reverse;
                else e_set.uv_type[1] = efem_sf::Normal;

                aij_p1 = g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Aij_part1, &e_set, this);

                aij_p2 = g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Aij_part2, &e_set, this);

                tmp_entry.real(aij_p1);
                tmp_entry.imag(-1.0*omega_*_mu0*edge_cndt*aij_p2);

                omp_entries[i*6 + j][e] = triplt_cd(elem_edge_index_[e][i], elem_edge_index_[e][j], tmp_entry);
            }
        }
    }

    for (i = 0; i < 36; i++)
    {
        kernel_entries.insert(kernel_entries.end(), omp_entries[i].begin(), omp_entries[i].end());
        omp_entries[i].clear();
    }

    A_PDE_.setFromTriplets(kernel_entries.begin(), kernel_entries.end());
    kernel_entries.clear();

    B_TE_.resize(edge_num_); B_TE_.setZero();
    B_TM_.resize(edge_num_); B_TM_.setZero();

    // TE
    for (int e = 0; e < elem_num_; ++e)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);
        get_element_conductivity(elem_markers_[e], edge_cndt_ref, true);
        edge_cndt -= edge_cndt_ref;

        if (edge_cndt != 0.0) // only calculate elements that have a residual conductivity
        {
            for (int i = 0; i < 6; i++)
            {
                e_set.uv_id[0] = i;
                if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
                else e_set.uv_type[0] = efem_sf::Normal;

                B_TE_[elem_edge_index_[e][i]] += _1i*omega_*_mu0*edge_cndt*
                    g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                        tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                        tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Bi_TE, &e_set, this);
            }
        }
    }

    // TM
    for (int e = 0; e < elem_num_; ++e)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);
        get_element_conductivity(elem_markers_[e], edge_cndt_ref, true);
        edge_cndt -= edge_cndt_ref;

        if (edge_cndt != 0.0) // only calculate elements that have a residual conductivity
        {
            for (int i = 0; i < 6; i++)
            {
                e_set.uv_id[0] = i;
                if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
                else e_set.uv_type[0] = efem_sf::Normal;

                B_TM_[elem_edge_index_[e][i]] += _1i*omega_*_mu0*edge_cndt*
                    g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                        tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                        tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Bi_TM, &e_set, this);
            }
        }
    }

    std::clog << "done\n";
    return;
}
